open Format
open Utils.Error
open Utils.Pervasive
open Utils.RandSelect
open Syntax
open Pointers
open Binding
open Core
open PrintExpr
open GPUtils




module Gene = struct
	type delete_status = NoDelete | Delete | NoSpec
	type gene = {
		expression : term;
		gene_id : int;
		
		mutable protection : delete_status; 
		mutable geneFitness : float;	
		mutable scoreOfBestCarrier : float;
		mutable selectionProb : float;
		mutable useInPopulation : int;
		mutable seedSelectionCount : int
	}
	let gene_to_expression g = Te(g.expression, None)
	let new_gene ?(selectionProb=0.0) ?(ds=NoSpec) te id = {
		expression = te;
		gene_id = id;
		
		protection = ds;
		seedSelectionCount = 0;
		scoreOfBestCarrier = 0.0;
		geneFitness = 0.0;
		selectionProb = selectionProb;
		useInPopulation = 0
	}
end 

open Gene
module GenePool = struct
	exception BlockConstructionFailure
	type gene_category_info = {
		mutable species_use : int;
		mutable total_fitness : float;
		mutable next_gene_id : int 
	}
	type gene_pool = {
		g : (ty, (gene_category_info * gene list)) Hashtbl.t;
		mutable ctx : context ref
	}
	
	let tyMem genes ty= 
		Hashtbl.fold 
		(fun tyElt _ acc -> acc || ty_eqv !(genes.ctx) tyElt ty) 
		genes.g false
	let tyFind genes ty = 
		Hashtbl.find genes.g (simplify_ty !(genes.ctx) ty)
	let teMem genes te= 
		let ty = typeof !(genes.ctx) te in
		if tyMem genes ty 
			then List.exists (fun t -> te_eqv !(genes.ctx) t.expression te) (snd (tyFind genes ty))
			else false
	let teFind genes te = 
	ListLabels.find (snd (tyFind genes (typeof !(genes.ctx) te)))
	~f:(fun t -> te_eqv !(genes.ctx) t.expression te)
	
	let get_genes genes = 
	Hashtbl.fold (fun _ (_, lst) accLst -> lst @ accLst)
	genes.g [ ]
	let get_term genes ty id = 
		let _, lst = tyFind genes ty in
		let gene = List.find (fun g-> g.gene_id = id) lst in
		gene.expression
	let get_types ?(pred =(fun _->true)) genes = 
	Hashtbl.fold 
	(fun ty _ acc -> if pred ty then ty::acc else acc) 
	genes.g []
	let get_terms ?(predTe =(fun _->true)) ?(predTy =(fun _->true)) genes  = 
	Hashtbl.fold 
		(fun ty (_, lst) acc -> 
			(if not (predTy ty) then [ ] else 
			List.filter (fun te -> predTe te)
			(List.map (fun gene -> gene.expression) lst)) @ acc) 
			genes.g [ ]	
		
	let reset_values genes = 
		Hashtbl.iter (fun _ (cinfo, lst) -> 
			cinfo.species_use <- 0;
			cinfo.total_fitness <- 0.0;
		ListLabels.iter lst
		~f:(fun g -> 
			g.geneFitness <- 0.0;
			g.selectionProb <- 0.0;
			g.useInPopulation <- 0;
			g.scoreOfBestCarrier <- 0.0)) genes.g		
						
				
	let rem_type genes ty = 
		let sTy = simplify_ty !(genes.ctx) ty in
		if tyMem genes sTy  then Hashtbl.remove genes.g sTy				
	let replace_type genes ty catInfo geneLst  = 
		let sTy = simplify_ty !(genes.ctx) ty in
		if tyMem genes sTy  
		then ( 
			let cInfo, lst = tyFind genes sTy in
			let catInfo  = match catInfo with 
				| Some cInfo -> cInfo 
				| None -> cInfo in
			let geneLst = match geneLst with 
				| Some newlst -> newlst
				| None -> lst in
			assert (catInfo.next_gene_id >= List.length geneLst);
			Hashtbl.replace genes.g sTy (catInfo,geneLst) )
		else (
			let geneLst = match geneLst with 
				| Some newlst -> newlst
				| None -> [ ] in	
			let catInfo  = match catInfo with 
				| Some cInfo -> cInfo 
				| None -> {species_use=0; total_fitness = 0.0; next_gene_id = List.length geneLst } in
			assert (catInfo.next_gene_id >= List.length geneLst);
			Hashtbl.add genes.g sTy (catInfo,geneLst))		
			
			
	let rec make_gene_pool ctx_ref = 
		let pool = {
			ctx = ref !ctx_ref;
			g = Hashtbl.create 200} in
	
		List.iter (fun (name, _ ) ->
		un (insert_expression pool (expr_of_name !ctx_ref name))) !ctx_ref;
		pool
		
	and insert_expression ?(ds = NoSpec) genes expr = 
		let inner_ins genes = function
	    |   Ty (ty,_) -> replace_type genes ty None None
	    |	Te (te, name) -> 
				match teMem genes te with 
				| false -> 
					let ty = typeof !(genes.ctx) te in
					replace_type genes ty None None;
					let cinfo, geneLst = tyFind genes ty in
					let id = cinfo.next_gene_id in 
					cinfo.next_gene_id <- cinfo.next_gene_id + 1;	
					let gene = new_gene ~ds te id in
					replace_type genes ty (Some cinfo) (Some (gene::geneLst)) 
				| true -> 
					(teFind genes te).protection <- ds
		in inner_ins genes expr

	let comp genes = 
	Hashtbl.fold 
	(fun _ (_ , lst) acc -> 
		ListLabels.fold_left ~init:acc lst
		~f:(fun acc g -> Pointers.complexity_of_expr (Te (g.expression, None)) + acc))
	genes.g 0 
	let gene_count genes = 
		Hashtbl.fold (fun _ (_, q) acc -> acc + List.length q) genes.g 0

	let is_gene_expression genes = function
	|Ty (ty, _)  -> tyMem genes ty 
	|Te (te, _) -> 
		let ctx = !(genes.ctx) in
		let ty = typeof ctx te in
		if not (tyMem genes ty) then false
		else 
			ListLabels.exists (snd (tyFind genes ty))
			~f:(fun t -> t.expression = te)  
end 

